/* Kernel module to match ruleid parameters. */

/* (C) 2013  <Leadsec>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.
 */

#include <linux/module.h>
#include <linux/skbuff.h>
#include <linux/seq_file.h>
#include <linux/jhash.h>
#include <linux/init.h>
#include <linux/net.h>
#include <linux/in.h>
#include <linux/ip.h>
#include <linux/netfilter/xt_esp.h>
#include <linux/netfilter/x_tables.h>
#include <linux/netfilter_ipv4/ip_tables.h>
#include <linux/netfilter_ipv6/ip6_tables.h>
#include <net/net_namespace.h>

#include <net/netfilter/nf_conntrack.h>
#include "xt_ruleid.h"

MODULE_LICENSE( "GPL" );
MODULE_DESCRIPTION( "x_tables rule Index match module" );
MODULE_ALIAS( "ipt_ruleid" );
MODULE_ALIAS( "ip6t_ruleid" );

#ifdef DEBUG_XT_RULEID
#define duprintf( format, args... ) printk( "%s %d: " format , __FUNCTION__, __LINE__, ## args )
#else
#define duprintf( format, args... )
#endif

static DEFINE_SPINLOCK( ruleid_lock );
static struct hlist_head *ruleid_head = NULL;

static ruleid_cfg_t cfg = {
	.ruleid_switch	= 1,
	.ruleid_info_max = 8096,
	.ruleid_info_amount = ATOMIC_INIT( 0 ),
	.memory_info = "\0",
	.total_alloc_size = 0,
	.total_alloc_realsize = 0,
	.getpages_size = 0,
	.getpages_realsize = 0,
	.getpages_number = 0,
	.mem_max_size = 0
};

/* ELF hash function */
static inline u_int32_t
ruleid_str_hash( const char *str )
{
        u_int32_t hash = 0 ;
        u_int32_t x = 0 ;

        while ( *str ) {
                hash = ( hash << 4 ) + ( *str ++ );
                if  ( (x = hash & 0xF0000000L ) != 0 ) {
                        hash ^= ( x >> 24 );
                        hash &=  ~ x;
                }
        }
        return ( hash & RULEID_HASH_MASK );
}

static void*
ruleid_kzalloc( size_t size, gfp_t flags )
{
        void *p = NULL;
        size_t realsize = 0;

        if ( size <= PAGE_SIZE ) {
                realsize = kmalloc_size( size );
                if ( cfg.total_alloc_realsize + realsize > cfg.mem_max_size && cfg.mem_max_size != 0 ) {
                        printk( "netobj alloc memory over %lu,"
                                "please increase by /proc/sys\n", (unsigned long)cfg.mem_max_size );
                        return NULL;
                }
                p = kmalloc( size, flags );
        } else {
                realsize = getpages_size( size );
                if ( cfg.total_alloc_realsize + realsize > cfg.mem_max_size && cfg.mem_max_size != 0 ) {
                        printk( "netobj alloc memory over %lu,"
                                "please increase by /proc/sys\n", (unsigned long)cfg.mem_max_size );
                        return NULL;
                }
                p = (void *)__get_free_pages( flags, get_order(size) );
        }

        if ( NULL == p )
                return NULL;
	memset( p, 0, size );

        if (size > PAGE_SIZE) {
                cfg.getpages_size += size;
                cfg.getpages_realsize += realsize;
                cfg.getpages_number++;
        }

        cfg.total_alloc_size += size;
        cfg.total_alloc_realsize += realsize;

        return p;
}

static void
ruleid_kfree( const void *p, size_t size )
{
        size_t realsize;

        if ( size <= PAGE_SIZE ) {
                kfree( p );
                realsize = kmalloc_size( size );
        } else {
                free_pages( (unsigned long)p, get_order(size) );
                realsize = getpages_size( size );
                cfg.getpages_size -= size;
                cfg.getpages_realsize -= realsize;
                cfg.getpages_number--;
        }

        cfg.total_alloc_size -= size;
        cfg.total_alloc_realsize -= realsize;
	
	return ;
}

static inline void
ruleid_info_get( ruleid_info_t *node )
{
	atomic_inc( &node->refcnt );
	return ;
}

static inline void
ruleid_info_put( ruleid_info_t *node )
{
	if ( atomic_dec_and_test(&node->refcnt) ) {
		hlist_del( &node->hnode );
		ruleid_kfree( node, node->size );
		atomic_dec( &cfg.ruleid_info_amount );
		/* TRAP */
	}
	return ;
}

static ruleid_info_t *
ruleid_find_name( const u_int8_t *name, u_int8_t create )
{
	u_int32_t hash = ruleid_str_hash( name );
	ruleid_info_t *node = NULL;
	struct hlist_node *p;
	size_t len = 0;

	if ( !name )
		return NULL;

	/* Find name */
	spin_lock_bh( &ruleid_lock );
	hlist_for_each_entry( node, p, &ruleid_head[hash], hnode ) {
		if ( node && strncmp( node->ruleid, name, RULEID_NAME_SIZE ) == 0 ) 
			goto unlock;
	}

	if ( !create ) 
		goto unlock;
	
	/* Check ruleid_info max value */
	atomic_inc( &cfg.ruleid_info_amount );
	
	if ( cfg.ruleid_info_max && 
	    unlikely(atomic_read(&cfg.ruleid_info_amount) > cfg.ruleid_info_max) ) {
		if ( net_ratelimit() ) 
			duprintf( "ruleid-match rulid_info amount is enough\n" );
		goto unlock;
	}

	/* Create new cache for ruleid_info */
	len = sizeof( ruleid_info_t );
	node = (ruleid_info_t*)ruleid_kzalloc( len, GFP_ATOMIC );
	if ( !node ) {
		atomic_dec( &cfg.ruleid_info_amount );
		goto unlock;
	}

	/* Initialize ruleid_info */
	node->size = len;
	strncpy( node->ruleid, name, RULEID_NAME_SIZE - 1 );
	atomic_set( &node->refcnt, 0 );
	spin_lock_init( &node->lock );
	hlist_add_head(&node->hnode, &ruleid_head[hash]);
	
unlock:
	spin_unlock_bh( &ruleid_lock );
	return node;
}

static inline int
ruleid_counter( const struct sk_buff *skb, const u_int8_t * name )
{
	struct nf_conn *ct = NULL;
	enum ip_conntrack_info ctinfo;
	ruleid_info_t *r_info = NULL;

	r_info = ruleid_find_name( name ,0 );
	if ( !r_info ) {
		if ( net_ratelimit() ) 
			duprintf( "Warning: Ruleid-match ruleinfo not find.\n" );
		return 0;
	}
	
	ct = nf_ct_get( skb, &ctinfo );
	if ( !ct ) {
		/* When packet isn't a session, such as the ah protocol, etc. But you shoule be count
 		   the number of packets and bytes.	 		
  		*/
		spin_lock_bh( &ruleid_lock );
		r_info->pcnt++;
		r_info->bcnt += skb->len;
		spin_unlock_bh( &ruleid_lock );
		return 0;
	}

	spin_lock_bh( &ruleid_lock );
	if ( !(nf_ct_is_confirmed(ct)) ) {
		r_info->scnt++;
		r_info->pcnt++;
		r_info->bcnt += skb->len;
	} else {
		r_info->pcnt++;
		r_info->bcnt += skb->len;
	}
	spin_unlock_bh( &ruleid_lock );

	return 0;
}

static bool
ruleid_mt( const struct sk_buff *skb, struct xt_action_param *par )
{
	const ruleid_match_t *arg = par->matchinfo;
	int ret = 0;

	if ( !cfg.ruleid_switch )	
		return true;

	ret = ruleid_counter( skb, arg->ruleid );
	if ( ret < 0 ) {
		duprintf( "Get conntrack failed.\n" );
		return true;
	}	
	return true;
}

static int
ruleid_mt_check( const struct xt_mtchk_param *par )
{
	ruleid_match_t *arg = par->matchinfo;
	ruleid_info_t *node = NULL;
	if ( !arg->ruleid )
		return EINVAL;

	node = ruleid_find_name( arg->ruleid, 1 );
	if ( !node ) {
		duprintf( "Create policy fauled because %s match.\n", par->match->name );
		return 1;
	}
	ruleid_info_get( node );

	return 0;
}

static void
ruleid_mt_destroy( const struct xt_mtdtor_param *par )
{
	ruleid_match_t *arg = par->matchinfo;
	struct hlist_node *p, *n;
	ruleid_info_t *node = NULL;
	u_int32_t hash = ruleid_str_hash( arg->ruleid );

	spin_lock_bh( &ruleid_lock );
	hlist_for_each_entry_safe( node, p, n, &ruleid_head[hash], hnode ) {
		if ( node && strncmp( node->ruleid, arg->ruleid, RULEID_NAME_SIZE ) == 0 ) {
			ruleid_info_put( node );
			spin_unlock_bh( &ruleid_lock );
			return ;
		}
	}
	spin_unlock_bh( &ruleid_lock );
	
	return ;
}

static int
do_ruleid_flush_cmd( u_int8_t *name )
{
	int i = 0;
	ruleid_info_t *node = NULL;
	struct hlist_node *p;

	if ( *name != '\0' ) {
		node = ruleid_find_name( name, 0 );
		if ( !node ) 
			return -1;
		node->pcnt = node->bcnt = node->scnt = 0;
		return 0;
	}

	spin_lock_bh( &ruleid_lock );
	for ( i = 0; i < RULEID_HASH_SIZE - 1; i++ ) 
		hlist_for_each_entry( node, p, &ruleid_head[i], hnode ) 
			node->pcnt = node->scnt = node->bcnt = 0;
	spin_unlock_bh( &ruleid_lock );

	return 0;
}

static int
do_ruleid_set_ctl( struct sock *sk, int cmd, void __user *user, unsigned int len )
{
	int ret;
	ruleid_info_t u_info;

	if ( ((u_int32_t)cmd) >= RULEID_SO_SET_MAX ) return -EINVAL;

	if ( sizeof(ruleid_info_t) != len ) {
		duprintf( "Invalid information length!\n" );
		return -ENOPROTOOPT;
	}

	ret = copy_from_user( &u_info, user, len );
	if ( 0 != ret ) {
		duprintf( "Can not copy information from userspace to kernel!\n" );
		return ret;
	}

	switch ( cmd ) {
	case RULEID_SO_SET_ENA:
		cfg.ruleid_switch = 1;
		break;

	case RULEID_SO_SET_DIS:
		cfg.ruleid_switch = 0;
		break;
		
	case RULEID_SO_FLUSH_CNT:
		ret = do_ruleid_flush_cmd( u_info.ruleid );
		break;

	default:
		break;
	}
	
	return ret;
}

static int
do_ruleid_get_ctl( struct sock *sk, int cmd, void __user *user, int *len )
{
	int ret;
	ruleid_info_t u_info, *k_info;

	if ( ((u_int32_t)cmd) >= RULEID_SO_GET_MAX ) return -EINVAL;

	if ( sizeof(ruleid_info_t) != *len ) {
		duprintf( "Invalid information length!\n" );
		return -ENOPROTOOPT;
	}

	ret = copy_from_user( &u_info, user, *len );
	if ( 0 != ret ) {
		duprintf( "Can not copy information from userspace to kernel!\n" );
		return ret;
	}

	k_info = ruleid_find_name( u_info.ruleid, 0 );
	if ( !k_info )
		return -1;

	spin_lock_bh( &ruleid_lock );
	switch ( cmd ) {
	case RULEID_SO_GET_CNT:
		u_info.bcnt = k_info->bcnt;
		u_info.pcnt = k_info->pcnt;
		u_info.scnt = k_info->scnt;
		break;
	
	default:
		duprintf( "do_ruleid_get_ctl cmd error.\n" );
		break;
	}
	spin_unlock_bh( &ruleid_lock );

	ret = copy_to_user( user, &u_info, *len );
	if ( 0 != ret ) {
		duprintf( "Can not copy information from kernel to userspace!\n" );
		return ret;
	}
		
	return ret;
}

static struct xt_match ruleid_mt_reg[] __read_mostly = {
	{
		.name		= "ruleid",
		.family		= NFPROTO_IPV4,
		.checkentry	= ruleid_mt_check,
		.match		= ruleid_mt,
		.destroy	= ruleid_mt_destroy,
		.matchsize	= sizeof( ruleid_match_t ),
		.me		= THIS_MODULE,
	},
	{
		.name		= "ruleid",
		.family		= NFPROTO_IPV6,
		.checkentry	= ruleid_mt_check,
		.match		= ruleid_mt,
		.destroy	= ruleid_mt_destroy,
		.matchsize	= sizeof( ruleid_match_t ),
		.me		= THIS_MODULE,
	},
};

/* Sockopt for config and get statis of rules */
static struct nf_sockopt_ops ruleid_sockopts = {
	.pf		= PF_INET,
	.set_optmin	= RULEID_BASE_CTL,
	.set_optmax	= RULEID_SO_SET_MAX + 1,
	.set		= do_ruleid_set_ctl,
	
	.get_optmin	= RULEID_BASE_CTL,
	.get_optmax	= RULEID_SO_GET_MAX + 1,
	.get		= do_ruleid_get_ctl,
	.owner		= THIS_MODULE,
};

#ifdef CONFIG_PROC_FS
static void *
ruleid_seq_start( struct seq_file *seq, loff_t *pos )
{
	if ( !(*pos) ) {
		seq_puts( seq, "RULE_ID     PACKET(S)        BYTE(S)        SESSION(S)     REFCNT(S) \n" );
	}
	
	while ( *pos < RULEID_HASH_SIZE && hlist_empty(&ruleid_head[*pos]) )
		(*pos)++;

	return *pos >= RULEID_HASH_SIZE ? NULL : pos;
}

static void *
ruleid_seq_next( struct seq_file *seq, void *v, loff_t *pos )
{
	(*pos)++;
	
	while ( *pos < RULEID_HASH_SIZE && hlist_empty(&ruleid_head[*pos]) )
		(*pos)++;

	return *pos >= RULEID_HASH_SIZE ? NULL : pos;
}

static void 
ruleid_seq_stop( struct seq_file *seq, void *v )
{
	return ;
}

static int ruleid_seq_show( struct seq_file *seq, void *v )
{
	struct hlist_node *p;
	ruleid_info_t *node = NULL;
	
	loff_t *pos = v;
	if ( NULL == v )
		return 0;
		
	spin_lock_bh( &ruleid_lock );	
	hlist_for_each_entry( node, p, &ruleid_head[*pos], hnode ) {
		seq_printf( seq, "%-12s%-17lld%-15lld%-16lld%d \n", node->ruleid, node->pcnt, 
				node->bcnt, node->scnt, atomic_read(&node->refcnt) );

	}
	spin_unlock_bh( &ruleid_lock );

	return 0;
}

static const struct seq_operations ruleid_seq_ops = {
	.start	= ruleid_seq_start,
	.next       = ruleid_seq_next,
	.stop       = ruleid_seq_stop,
	.show      = ruleid_seq_show,
};

static int ruleid_seq_open( struct inode *inode, struct file *file )
{
	return seq_open( file, &ruleid_seq_ops );
}

static const struct file_operations ruleid_file_fops = {
	.owner          = THIS_MODULE,
	.open		= ruleid_seq_open,
	.read         	= seq_read,
	.llseek         = seq_lseek,
	.release        = seq_release,
};

struct proc_dir_entry *proc_ruleid = NULL;
static int __net_init 
ruleid_net_init( struct net *net )
{
	struct proc_dir_entry *entry = NULL;
	
	proc_ruleid = proc_mkdir( "ruleid", proc_leadsec );
	if ( !proc_ruleid )
		return -ENOMEM;

	entry = proc_create_data( "rule_stat", S_IRUGO, proc_ruleid, &ruleid_file_fops, NULL );
	if ( !entry )
		return -ENOMEM;

	proc_symlink( "memory_info", proc_ruleid, PROC_PATH"/memory_info" );
	proc_symlink( "memory_max", proc_ruleid, PROC_PATH"/memory_max" );
	proc_symlink( "memory_use", proc_ruleid, PROC_PATH"/memory_use" );
	proc_symlink( "ruleid_info_amount", proc_ruleid, PROC_PATH"/ruleid_info_amount" );
	proc_symlink( "ruleid_info_max", proc_ruleid, PROC_PATH"/ruleid_info_max" );

	return 0;
}

static void __net_exit 
ruleid_net_exit( struct net *net )
{
	remove_proc_entry( "memory_info", proc_ruleid );
	remove_proc_entry( "memory_max", proc_ruleid );
	remove_proc_entry( "memory_use", proc_ruleid );
	remove_proc_entry( "ruleid_info_amount", proc_ruleid );
	remove_proc_entry( "ruleid_info_max", proc_ruleid );

	remove_proc_entry( "rule_stat", proc_ruleid );
	remove_proc_entry( "ruleid", proc_leadsec );
	return ;
}

static struct pernet_operations ruleid_net_ops = {
	.init = ruleid_net_init,
	.exit = ruleid_net_exit,
};

static int
ruleid_read_proc_init( void )
{
	return register_pernet_subsys( &ruleid_net_ops );
}

static void
ruleid_read_proc_exit( void )
{
	unregister_pernet_subsys( &ruleid_net_ops );
}

#else	/* CONFIG_PROC_FS */

static int
ruleid_read_proc_init( void )
{
	return 0;
}

static void
ruleid_read_proc_exit( void )
{
	return ;
}
#endif	/* CONFIG_PROC_FS */

#ifdef CONFIG_SYSCTL
struct ctl_table_header *ruleid_sysctl_header = NULL;

static int ruleid_sysctl_mem_info(ctl_table *ctl, int write,
                           void __user *buffer, size_t *lenp, loff_t *ppos)
{
	int pos = 0, len = 0;
	char *info = cfg.memory_info;
	const int max_size = sizeof( cfg.memory_info );

	if (!*lenp || (*ppos && !write)) {
		*lenp = 0;
		return 0;
	}

	len = scnprintf( info + pos, max_size - pos, 
		"Max memory for ruleid module:\t%d", cfg.mem_max_size );
	if ( !len  )
		goto done;
	pos += len;

	len = scnprintf( info + pos, max_size - pos, 
		"\nTotal memory for allocated:\t%d", cfg.total_alloc_realsize );
	if ( !len )
		goto done;
	pos += len;

	len = scnprintf( info + pos, max_size - pos,
		"\nTotal memory for use:\t%d", cfg.total_alloc_size );
	if ( !len )
		goto done;
	pos += len;

	len = scnprintf( info + pos, max_size - pos,
		"\nTotal getpages number:\t%d\n", cfg.getpages_number );
	if ( !len )
		goto done;

doit:
	return proc_dostring( ctl, write, buffer, lenp, ppos );

done:
	duprintf( "Too small!!!\n" );
	goto doit;
}

static ctl_table ruleid_sysctl_table[] = {
	{
		.procname		= "ruleid_info_max",
		.data			= &cfg.ruleid_info_max,
		.maxlen			= sizeof( cfg.ruleid_info_max ),
		.mode			= 0644,
		.proc_handler	= &proc_dointvec,
		
	},
	{
		.procname		= "ruleid_info_amount",
		.data			= &cfg.ruleid_info_amount,
		.maxlen			= sizeof( cfg.ruleid_info_amount ),
		.mode			= 0444,
		.proc_handler		= &proc_dointvec,
		
	},
	{
		.procname		= "memory_info",
		.data			= &cfg.memory_info,
		.maxlen			= sizeof( cfg.memory_info ),
		.mode			= 0444,
		.proc_handler	= &ruleid_sysctl_mem_info,
	},
	{
		.procname		= "memory_use",
		.data			= &cfg.total_alloc_size,
		.maxlen			= sizeof( cfg.total_alloc_size ),
		.mode			= 0444,
		.proc_handler	= &proc_dointvec,
	},
	{
		.procname		= "memory_max",
		.data			= &cfg.mem_max_size,
		.maxlen			= sizeof( cfg.mem_max_size ),
		.mode			= 0644,
		.proc_handler	= 	&proc_dointvec,
	},

	{ }
};

static ctl_table ruleid_child_table[] = {
	{
		.procname	= "ruleid",
		.mode		= 0555,
		.child		= ruleid_sysctl_table,
	},
	{ }
};

static ctl_table ruleid_leadsec_table[] = {
	{
		.procname	= "leadsec",
		.mode		= 0555,
		.child		= ruleid_child_table,
	},
	{ }
};

static int ruleid_sysctl_init( void )
{
	ruleid_sysctl_header = register_sysctl_table( ruleid_leadsec_table );
	if ( NULL == ruleid_sysctl_header )
		return -ENOMEM;

	return 0;
}

static void ruleid_sysctl_fint( void )
{
	unregister_sysctl_table( ruleid_sysctl_header );
	return ;
}

#else	/* CONFIG_SYSCTL */

static int ruleid_sysctl_init( void )
{
	return 0;
}


static void ruleid_sysctl_fint( void )
{
	return ;
}

#endif	/* CONFIG_SYSCTL */

static int ruleid_proc_init( void )
{
	int ret = -1;

	ret = ruleid_sysctl_init();
	if ( 0!= ret )
		return -ENOMEM;

	ret = ruleid_read_proc_init();
	if ( 0 != ret )
		goto unreg_sysctl;

	return ret;

unreg_sysctl:
	ruleid_sysctl_fint();
	return ret;
}

static void ruleid_proc_exit( void )
{
	ruleid_sysctl_fint();
	ruleid_read_proc_exit();
	return ;
}

static int __init xt_ruleid_init( void )
{
	int ret = 0;
	int i = 0;

	ruleid_head = kzalloc( RULEID_HASH_SIZE * sizeof(struct hlist_head), GFP_KERNEL );
	if ( !ruleid_head ) {
		duprintf( "Kzalloc memory for stat info failed.\n" );
		return -ENOMEM;
	}

	for ( i =0; i < RULEID_HASH_SIZE; i++ ) {
		INIT_HLIST_HEAD( &ruleid_head[i] );
	}

	ret = nf_register_sockopt( &ruleid_sockopts );
	if ( 0 != ret ) {
		duprintf( "Reigster sockopt ruleid_sockopts failed.\n" );
		goto free;
			
	}

	ret = xt_register_matches( ruleid_mt_reg, ARRAY_SIZE(ruleid_mt_reg) );
	if ( 0 != ret ) {
		duprintf( "Register iptables match failed.\n" );
		goto unregister_sockopt;
	}

	ret = ruleid_proc_init();
	if ( 0 != ret ) {
		duprintf( "Register proc file system failed.\n" );
		goto unregister_match;
	}
	duprintf( "xt_ruleid already load.\n" );
	
	return ret;

unregister_match:
	xt_unregister_matches( ruleid_mt_reg, ARRAY_SIZE(ruleid_mt_reg) );	
unregister_sockopt:
	nf_unregister_sockopt( &ruleid_sockopts );
free:
	kfree( ruleid_head );

	return -1;
}

static void __exit xt_ruleid_cleanup( void )
{
	nf_unregister_sockopt( &ruleid_sockopts );
	xt_unregister_matches( ruleid_mt_reg, ARRAY_SIZE(ruleid_mt_reg) );
	ruleid_proc_exit();
	kfree( ruleid_head );
	duprintf( "xt_ruleid already exit.\n" );
	return ;
}

module_init(xt_ruleid_init);
module_exit(xt_ruleid_cleanup);
